import os
import time
import json
import psutil
import ctypes
import paho.mqtt.client as mqtt

# -------------------------------
# 调试开关
# -------------------------------
debug = True  # 将此设置为 False 可以关闭所有 print 输出


def debug_print(message, *args):
    if debug:
        if args:
            print(message, *args)
        else:
            print(message)


# -------------------------------
# 读取 JSON 配置文件
# -------------------------------
with open('config.json', 'r', encoding='utf-8') as f:
    config = json.load(f)

broker = config["broker"]
username = config["username"]
password = config["password"]
port = config["port"]
client_id = config["client_id"]

# -------------------------------
# MQTT 主题定义（嵌入 client_id 使每台电脑主题唯一）
# -------------------------------
status_topic = f'home/computer/{client_id}/status/uptime'
shutdown_topic = f'home/computer/{client_id}/command/shutdown'
lock_topic = f'home/computer/{client_id}/command/lock'
reboot_topic = f'home/computer/{client_id}/command/reboot'
lock_status_topic = f'home/computer/{client_id}/status/lock'
reboot_status_topic = f'home/computer/{client_id}/status/reboot'

# 各状态数据主题
cpu_topic = f'home/computer/{client_id}/status/cpu'
memory_topic = f'home/computer/{client_id}/status/memory'
upload_speed_topic = f'home/computer/{client_id}/status/upload_speed'
download_speed_topic = f'home/computer/{client_id}/status/download_speed'

# -------------------------------
# MQTT 客户端初始化与回调函数
# -------------------------------
client = mqtt.Client(client_id)
client.username_pw_set(username, password)


def publish_discovery_config():
    """
    发布 Home Assistant 自动发现配置
    """
    device_info = {
        "identifiers": [client_id],
        "name": "HW笔记本",
        "manufacturer": "HuaWwei",
        "model": "PC-Model",
        "sw_version": "1.0"
    }

    # CPU 使用率传感器
    cpu_payload = {
        "unique_id": f"{client_id}_cpu_usage",
        "name": "CPU 占用",
        "state_topic": cpu_topic,
        "unit_of_measurement": "%",
        "device_class": "power",
        "value_template": "{{ value }}",
        "device": device_info
    }
    client.publish(f'homeassistant/sensor/{client_id}/cpu/config', json.dumps(cpu_payload), retain=True, qos=1)

    # 内存使用率传感器
    memory_payload = {
        "unique_id": f"{client_id}_memory_usage",
        "name": "内存使用",
        "state_topic": memory_topic,
        "unit_of_measurement": "%",
        "value_template": "{{ value }}",
        "device": device_info
    }
    client.publish(f'homeassistant/sensor/{client_id}/memory/config', json.dumps(memory_payload), retain=True, qos=1)

    # 上传速度传感器
    upload_payload = {
        "unique_id": f"{client_id}_upload_speed",
        "name": "上传速度",
        "state_topic": upload_speed_topic,
        "unit_of_measurement": "Mbps",
        "value_template": "{{ value }}",
        "device": device_info
    }
    client.publish(f'homeassistant/sensor/{client_id}/upload_speed/config', json.dumps(upload_payload), retain=True,
                   qos=1)

    # 下载速度传感器
    download_payload = {
        "unique_id": f"{client_id}_download_speed",
        "name": "下载速度",
        "state_topic": download_speed_topic,
        "unit_of_measurement": "Mbps",
        "value_template": "{{ value }}",
        "device": device_info
    }
    client.publish(f'homeassistant/sensor/{client_id}/download_speed/config', json.dumps(download_payload), retain=True,
                   qos=1)

    # 关机开关配置
    shutdown_payload = {
        "unique_id": f"{client_id}_shutdown_switch",
        "name": "开机状态",
        "state_topic": status_topic,
        "command_topic": shutdown_topic,
        "payload_on": "on",
        "payload_off": "off",
        "state_on": "on",
        "state_off": "off",
        "device": device_info
    }
    client.publish(f'homeassistant/switch/{client_id}/shutdown/config', json.dumps(shutdown_payload), retain=True,
                   qos=1)

    # 锁屏开关配置
    lock_payload = {
        "unique_id": f"{client_id}_lock_switch",
        "name": "锁屏",
        "state_topic": lock_status_topic,
        "command_topic": lock_topic,
        "payload_on": "lock",
        "payload_off": "unlock",
        "state_on": "locked",
        "state_off": "unlocked",
        "device": device_info
    }
    client.publish(f'homeassistant/switch/{client_id}/lock/config', json.dumps(lock_payload), retain=True, qos=1)

    # 重启开关配置
    reboot_payload = {
        "unique_id": f"{client_id}_reboot_switch",
        "name": "重启",
        "state_topic": reboot_status_topic,
        "command_topic": reboot_topic,
        "payload_on": "reboot",
        "payload_off": "reboot",
        "state_on": "rebooting",
        "state_off": "running",
        "device": device_info
    }
    client.publish(f'homeassistant/switch/{client_id}/reboot/config', json.dumps(reboot_payload), retain=True, qos=1)

    debug_print("已发布所有传感器和开关配置到 Home Assistant")


def on_connect(client, userdata, flags, rc):
    if rc == 0:
        debug_print("MQTT 连接成功")
        # 连接成功后，订阅命令主题（包含 client_id），使用 QoS=1
        client.subscribe(shutdown_topic, qos=1)
        client.subscribe(lock_topic, qos=1)
        client.subscribe(reboot_topic, qos=1)
        # 重新发布 Home Assistant 自动发现配置
        publish_discovery_config()
    else:
        debug_print("MQTT 连接失败，返回码：", rc)


def on_disconnect(client, userdata, rc):
    debug_print("MQTT 连接断开，返回码：", rc)
    # 如果不是正常断开（rc != 0），则尝试重连
    if rc != 0:
        while True:
            try:
                client.reconnect()
                debug_print("MQTT 重连成功")
                break
            except Exception as e:
                debug_print("MQTT 重连失败，5秒后重试：", e)
                time.sleep(5)


def on_message(client, userdata, msg):
    try:
        payload = msg.payload.decode()
    except Exception as e:
        debug_print("解码消息出错：", e)
        return

    debug_print(f"接收到消息：主题={msg.topic}，内容={payload}")

    # 关机命令处理
    if msg.topic == shutdown_topic and payload == "off":
        debug_print("收到关机命令，准备关机...")
        send_shutdown_signal()

    # 锁屏和解锁命令处理
    elif msg.topic == lock_topic:
        if payload == "lock":
            debug_print("收到锁屏命令，准备锁屏...")
            lock_screen()
        elif payload == "unlock":
            debug_print("收到解锁命令，但无法自动解锁，请手动解锁。")
            client.publish(lock_status_topic, "unlocked", retain=True, qos=1)
        else:
            debug_print(f"未知锁屏命令：{payload}")

    # 重启命令处理
    elif msg.topic == reboot_topic:
        if payload == "reboot":
            debug_print("收到重启命令，准备重启...")
            restart_computer()
        else:
            debug_print(f"收到未知重启命令：{payload}，将忽略。")
    else:
        debug_print(f"未知命令：{payload}")


client.on_connect = on_connect
client.on_disconnect = on_disconnect
client.on_message = on_message

# 建立连接，设置较短的 keepalive 值（例如30秒）
try:
    client.connect(broker, port, keepalive=30)
except Exception as e:
    debug_print("初始连接失败：", e)
    exit(1)

client.loop_start()


# -------------------------------
# 单位换算函数（网络速度）
# -------------------------------
def convert_speed(bytes_value):
    """
    将字节每秒转换为兆比特每秒（Mbps）
    1 字节 = 8 比特
    1 Mbps = 1024 * 1024 比特/秒
    """
    bits_per_second = bytes_value * 8  # 转换为比特每秒
    mbps = bits_per_second / (1024 * 1024)  # 转换为兆比特每秒
    return round(mbps, 2)  # 返回数值，不包含单位


# -------------------------------
# 网络数据初始化
# -------------------------------
net_io = psutil.net_io_counters()
bytes_sent = net_io.bytes_sent
bytes_recv = net_io.bytes_recv
last_time = time.time()

# 获取所有网络接口
net_if_stats = psutil.net_if_stats()
net_if_io = psutil.net_io_counters(pernic=True)
active_interfaces = {}

# 初始化每个活动接口的数据
for iface, stats in net_if_stats.items():
    if stats.isup:  # 只处理已启用的接口
        active_interfaces[iface] = {
            'bytes_sent': net_if_io[iface].bytes_sent,
            'bytes_recv': net_if_io[iface].bytes_recv,
            'last_time': time.time()
        }


# -------------------------------
# 状态发布函数
# -------------------------------
def publish_system_status():
    client.publish(status_topic, "on", retain=True, qos=1)
    debug_print("系统状态: on 已发布到 Home Assistant")


# -------------------------------
# 新增功能：清零各项状态
# -------------------------------
def clear_system_status():
    """
    在关机前将 CPU、内存、上传和下载数据清零，并发布到相应 MQTT 主题
    """
    client.publish(cpu_topic, 0, retain=True, qos=1)
    client.publish(memory_topic, 0, retain=True, qos=1)
    client.publish(upload_speed_topic, 0, retain=True, qos=1)
    client.publish(download_speed_topic, 0, retain=True, qos=1)
    debug_print("已将 CPU、内存、上传和下载数据清零，并发布到 Home Assistant")


# -------------------------------
# 关机操作
# -------------------------------
def send_shutdown_signal():
    debug_print("发送关机状态到 Home Assistant...")
    # 在关机前先清零各项状态
    clear_system_status()
    client.publish(status_topic, "off", retain=True, qos=1)
    debug_print("正在关机...")
    os.system("shutdown /s /t 1")
    client.loop_stop()


# -------------------------------
# 锁屏操作
# -------------------------------
def lock_screen():
    debug_print("执行锁屏操作...")
    os.system("rundll32.exe user32.dll,LockWorkStation")
    # 锁屏后状态由主循环检测更新


# -------------------------------
# 重启操作
# -------------------------------
def restart_computer():
    debug_print("发送重启状态到 Home Assistant...")
    client.publish(reboot_status_topic, "rebooting", retain=True, qos=1)
    debug_print("正在重启...")
    os.system("shutdown /r /t 1")


# -------------------------------
# 主循环
# -------------------------------
while True:
    # 检查是否存在关机信号文件（例如 D:\signal.txt），存在则触发关机
    signal_file_path = 'D:\\signal.txt'
    if os.path.exists(signal_file_path):
        send_shutdown_signal()
        try:
            os.remove(signal_file_path)
        except Exception as e:
            debug_print("删除信号文件出错：", e)
        break

    # 获取 CPU 和内存使用率（注意 interval=1 会阻塞 1 秒）
    cpu_usage = psutil.cpu_percent(interval=1)
    memory_usage = psutil.virtual_memory().percent

    # 获取网络上传与下载速度
    current_time = time.time()
    total_upload_speed = 0
    total_download_speed = 0

    # 获取当前网络接口数据
    current_net_io = psutil.net_io_counters(pernic=True)

    # 检查是否有新的活动接口
    for iface, stats in net_if_stats.items():
        if stats.isup and iface not in active_interfaces:
            active_interfaces[iface] = {
                'bytes_sent': current_net_io[iface].bytes_sent,
                'bytes_recv': current_net_io[iface].bytes_recv,
                'last_time': current_time
            }
            debug_print(f"发现新的活动网络接口: {iface}")

    # 计算每个活动接口的速度
    for iface, data in list(active_interfaces.items()):
        if iface in current_net_io:
            time_diff = current_time - data['last_time']
            if time_diff > 0:
                upload_speed = (current_net_io[iface].bytes_sent - data['bytes_sent']) / time_diff
                download_speed = (current_net_io[iface].bytes_recv - data['bytes_recv']) / time_diff

                # 更新接口数据
                active_interfaces[iface]['bytes_sent'] = current_net_io[iface].bytes_sent
                active_interfaces[iface]['bytes_recv'] = current_net_io[iface].bytes_recv
                active_interfaces[iface]['last_time'] = current_time

                # 累加总速度
                total_upload_speed += upload_speed
                total_download_speed += download_speed

                # 打印每个接口的速度
                debug_print(
                    f"接口 {iface}: 上传 {convert_speed(upload_speed)} Mbps, 下载 {convert_speed(download_speed)} Mbps")
        else:
            # 如果接口不再存在，从活动列表中移除
            del active_interfaces[iface]
            debug_print(f"移除不存在的网络接口: {iface}")

    # 转换总速度为 Mbps
    upload_speed_value = convert_speed(total_upload_speed)
    download_speed_value = convert_speed(total_download_speed)

    debug_print(f"CPU: {cpu_usage}%, 内存: {memory_usage}%")
    debug_print(f"总上传速度: {upload_speed_value} Mbps, 总下载速度: {download_speed_value} Mbps")

    # 发布各项状态到 MQTT（均使用 retain=True 和 QoS=1）
    client.publish(cpu_topic, cpu_usage, retain=True, qos=1)
    client.publish(memory_topic, memory_usage, retain=True, qos=1)
    client.publish(upload_speed_topic, str(upload_speed_value), retain=True, qos=1)
    client.publish(download_speed_topic, str(download_speed_value), retain=True, qos=1)

    # 发布系统状态
    publish_system_status()

    # 每2秒循环一次
    time.sleep(2)
